# 方法一的改进
class ListNode:
    def __init__(self, x, key):
        self.val = x
        self.key = key
        self.next = None
        self.prev = None
        self.count = 1


class LFUCache:

    def __init__(self, capacity: int):
        self.map = dict()  # { key: node }
        self.counter = dict()  # { 1: (first, last) }
        self.cap = capacity
        self.size = 0

    def get(self, key: int) -> int:
        if self.size == 0:
            return -1
        if key in self.map:
            node = self.map[key]
            self.incrKeyUsedCount(node)
            return node.val
        else:
            return -1

    def put(self, key: int, value: int) -> None:
        if key in self.map:
            node = self.map[key]
            node.val = value
            self.incrKeyUsedCount(node)
            return

        if self.size < self.cap:
            node = ListNode(value, key)
            self.map[key] = node
            self.addCount(node, 1)
            self.size += 1
        else:
            self.removeLastKey()
            self.size -= 1
            self.put(key, value)

    def addCount(self, node: ListNode, count: int) -> None:
        node.count = count
        if count in self.counter:
            [first, last] = self.counter[count]
            node.next = first
            first.prev = node
            self.counter[count] = [node, last]
        else:
            self.counter[count] = [node, node]

    def removeLastKey(self) -> None:
        if self.size == 0:
            return

        leastUsedCount = sorted(self.counter.keys())[0]
        [first, last] = self.counter[leastUsedCount]
        if first == last:
            self.counter.pop(leastUsedCount)
        else:
            prevlast = last.prev
            prevlast.next = None
            self.counter[leastUsedCount] = [first, prevlast]
        self.map.pop(last.key)

    def incrKeyUsedCount(self, node: ListNode) -> None:
        [first, last] = self.counter[node.count]

        if first == last:
            self.counter.pop(node.count)
        else:
            prev = node.prev
            next = node.next
            if node == first:
                next.prev = None
                node.next = None
                self.counter[node.count] = [next, last]
            elif node == last:
                prev.next = None
                node.prev = None
                self.counter[node.count] = [first, prev]
            else:
                next.prev = node.prev
                prev.next = node.next
                node.next = None
                node.prev = None

        node.count +=1

        self.addCount(node, node.count)


cache = LFUCache(2)
cache.put(1, 1)
cache.put(2, 2)
print(cache.get(1))   # 返回 1
cache.put(3, 3)      # 去除 key 2
print(cache.get(2))   # 返回 -1 (未找到key 2)
print(cache.get(3))   # 返回 3
cache.put(4, 4)      # 去除 key 1
print(cache.get(1))   # 返回 -1 (未找到 key 1)
print(cache.get(3))   # 返回 3
print(cache.get(4))   # 返回 4


c2 = LFUCache(2)
c2.put(3, 1)
c2.put(2, 1)
c2.put(2, 2)
c2.put(4, 4)
print(c2.get(2))  # 返回 2

c3 = LFUCache(2)
c3.put(2, 1)
c3.put(1, 1)
c3.put(2, 3)
c3.put(4, 1)
print(c3.get(1))  # -1
print(c3.get(2))  # 3
